#ifndef _TRT_H_
#define _TRT_H_

#include <string>
#include <vector>
#include <iostream>
#include <fstream>

#include "NvInfer.h"
#include "NvOnnxParser.h"
#include "NvInferRuntimeCommon.h"
#include <cuda_runtime_api.h>

using namespace nvinfer1;
using namespace std;

class Logger : public ILogger {
    void log(ILogger::Severity severity, const char* msg) noexcept {
        if (severity != Severity::kINFO)
            std::cout << msg << std::endl;
    }
};


class Trt{
public:
    Trt(){};

    ~Trt(){

    };
    
    void infer();

    void create_engine(const std::string& onnx_file);

    void print_info();

    void** get_io();

private:
    int serialize(const std::string& onnx_file, const std::string& engine_file);

    int deserialize(const std::string& engine_file);

    void allocate_buffer();

private:
    void* m_buffers_gpu_[4]; // 0: input_buffer 1~3: output_buffer

    void* m_buffers_cpu_[4]; // 0: input_buffer 1~3: output_buffer

    Logger m_glogger_;
    nvinfer1::IExecutionContext* m_context_;
    nvinfer1::ICudaEngine* m_engine_ = nullptr;
};

#endif